"""OCR处理模块，封装PaddleOCR功能"""
import os
import time
from typing import List, Tuple, Optional
from paddleocr import PaddleOCR
from utils import logger

class OCRProcessor:
    """OCR处理类"""
    
    def __init__(self, use_gpu: bool = False):
        """
        初始化OCR引擎
        
        :param use_gpu: 是否使用GPU加速
        """
        self.ocr = PaddleOCR(
            lang="ch",
            use_angle_cls=True,
            use_gpu=use_gpu,
            det_db_box_thresh=0.1,
            show_log=False
        )
        self.log = logger.get_logger("OCRProcessor")
        self.log.info("OCR引擎初始化完成")
    
    def process_image(self, img_path: str) -> Tuple[List[str], List[List[Tuple[float, float]]], List[float]]:
        """
        对图像进行OCR识别
        
        :param img_path: 图像文件路径
        :return: (文本列表, 坐标框列表, 置信度列表)
        """
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"图像文件不存在: {img_path}")
        
        try:
            start_time = time.time()
            raw_result = self.ocr.ocr(img_path, det=True, rec=True, cls=True)
            
            if not raw_result or not raw_result[0]:
                return [], [], []
                
            texts, boxes, scores = [], [], []
            
            for detection in raw_result[0]:
                if not detection or len(detection) < 2:
                    continue
                    
                # 提取位置信息和文本/置信度
                box = detection[0]
                text, confidence = detection[1]
                
                texts.append(text)
                boxes.append(box)
                scores.append(float(confidence))
            
            elapsed_time = time.time() - start_time
            self.log.info(f"识别完成: {len(texts)}个文本, 耗时{elapsed_time:.2f}秒")
            return texts, boxes, scores
        
        except Exception as e:
            self.log.error(f"OCR处理异常: {str(e)}")
            return [], [], []